#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab

example_array = np.array([8902, 14834, 7832, 14356, 25817, 20284, 23786,
                          8045, 16342, 21132, 11860, 8585, 12450, 21837,
                          24210, 25297, 23652, 20222, 16525, 19416, 22314,
                          29174, 19409, 6713, 10768, 12266, 28711, 23564,
                          23215, 17248, 20730, 17453, 10959, 14308, 12004,
                          7373, 8352, 11177, 15215, 24394, 19529, 24203,
                          17427, 15545, 11108, 12838, 9589, 5086, 9320,
                          8810, 16522, 22810, 26158, 22763, 19217, 12515,
                          13147, 11390, 13768, 6254, 6990, 7078, 19948,
                          25805, 18197, 26659, 14797, 15518, 31978, 17488,
                          9171, 5685, 4606, 4999, 13842, 17479, 12992,
                          24680, 19376, 15313, 19262, 9621, 7844, 10334,
                          10869, 12423, 22781, 25657
                          ], dtype=float)


def moving_avg(X, k=3, t=False):
    '''
    Returns an array with moving averages of given array with a centered
    lag of k.

    X - array of data.
    k - size of convolve window; default 3
    t - Boolean to LN transform, default False
    '''
    if t == False:
        pass
    else:
        X = np.log(X)
    weightings = np.repeat(1.0, k) / k
    return np.convolve(X, weightings)[k - 1: - (k - 1)]


#def sacf(X, k=1):
#    """
#    Computes the sample autocorrelation function coeffficient.
#    for given lag k. Default k of 1.
#    """
#    if k == 0:
#       return 1.0
#    flen = float(len(X))
#    ybar = float(sum([x for x in X])) / flen
#    D = sum([(x-xbar)**2 for x in X])
#    N = sum([(x-xbar)* (xtpk -xbar) for (x, xtpk) in zip(X[:-k],X[k:])])
#    return N/D


def autocorr(x, lag=1):
    """Returns the autocorrelation x."""
    x = np.squeeze(np.asarray(x))
    mu = x.mean()
    s = x.std()
    ac = ((x[: - lag] - mu) * (x[lag:] - mu)).sum() / (s ** 2) / (len(x) - lag)
    return ac

a = moving_avg(example_array)
b = autocorr(example_array)

q = example_array
plt.acorr(q, maxlags=None, detrend=mlab.detrend_linear, color='g', lw='2')
plt.grid(True)
plt.axhline(0, color='black', lw=2,)


print a
print b
print example_array
plt.show()
